from langgraph.graph import StateGraph, END
from langgraph.checkpoint.memory import MemorySaver
from langgraph.prebuilt import ToolNode

from new_editor.utils.tools import web_search_api
from state import EditorState
from title_agent import generate_title
from search_agent import search_resource
from outline_agent import generate_outline
from write_agent import generate_content
from check_agent import generate_check, check_to_write_agent
from comment_agent import generate_comment, comment_to_write_agent
from summary_agent import generate_summary

tools = [web_search_api]
tool_node = ToolNode(tools)

workflow = StateGraph(state_schema=EditorState)
workflow.add_node("title_agent", generate_title)
workflow.add_node("search_agent", search_resource)
workflow.add_node("outline_agent", generate_outline)
workflow.add_node("write_agent", generate_content)
workflow.add_node("check_agent", generate_check)
workflow.add_node("comment_agent", generate_comment)
workflow.add_node("summary_agent", generate_summary)
workflow.add_node("tool_node", tool_node)

workflow.set_entry_point("title_agent")

workflow.add_edge("title_agent", "search_agent")
workflow.add_edge("search_agent", "outline_agent")
workflow.add_edge("outline_agent", "write_agent")
workflow.add_edge("write_agent", "check_agent")
workflow.add_edge("check_agent", "comment_agent")
workflow.add_edge("comment_agent", END)

workflow.add_conditional_edges("check_agent",
                               check_to_write_agent,
                               {"comment_agent": "comment_agent", "write_agent": "write_agent"})

workflow.add_conditional_edges("comment_agent",
                               comment_to_write_agent,
                               {"summary_agent": "summary_agent", "write_agent": "write_agent"})

saver = MemorySaver()
app = workflow.compile(checkpointer=saver)

response = app.invoke({"title": "中国的近代史", "paragraphs": []}, config={"configurable": {"thread_id": 34}})
print(response)
